{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Copyright 2020 The OpenFermion Developers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FQE vs OpenFermion vs Cirq: Quadratic Hamiltonian Evolution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://quantumai.google/openfermion/fqe/tutorials/fqe_vs_openfermion_quadratic_hamiltonians\"><img src=\"https://quantumai.google/site-assets/images/buttons/quantumai_logo_1x.png\" />View on QuantumAI</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/quantumlib/OpenFermion/blob/master/docs/fqe/tutorials/fqe_vs_openfermion_quadratic_hamiltonians.ipynb\"><img src=\"https://quantumai.google/site-assets/images/buttons/colab_logo_1x.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/quantumlib/OpenFermion/blob/master/docs/fqe/tutorials/fqe_vs_openfermion_quadratic_hamiltonians.ipynb\"><img src=\"https://quantumai.google/site-assets/images/buttons/github_logo_1x.png\" />View source on GitHub</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://storage.googleapis.com/tensorflow_docs/OpenFermion/docs/fqe/tutorials/fqe_vs_openfermion_quadratic_hamiltonians.ipynb\"><img src=\"https://quantumai.google/site-assets/images/buttons/download_icon_1x.png\" />Download notebook</a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrates how evolve a state under a quadratic Hamiltonian with FQE, Cirq, and OpenFermion.\n",
    "\n",
    "We first generate a random quadratic hamiltonian using OpenFermion.  Then we set up a full $2^{n} \\times 2^{n}$ operator and a computational basis state corresponding to a Slater determinant with the zero and one indexed alpha spin-orbitals occupied.  We then compare the full space time evolutions to the LU-decomposition evolution implemented in FQE. Lastly, we using the `optimal_givens_decomposition` to implement the quadratic Hamiltonian time evolution in Cirq.  The 1-RDMs are compared for correctness.  We purposefully picked a high spin state so the spin-summed 1-RDM from FQE is equivalent to the spinless 1-RDM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import fqe\n",
    "except ImportError:\n",
    "    !pip install fqe --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Union\n",
    "from itertools import product\n",
    "import fqe\n",
    "from fqe.hamiltonians.restricted_hamiltonian import RestrictedHamiltonian\n",
    "\n",
    "import numpy as np\n",
    "import cirq\n",
    "\n",
    "import openfermion as of\n",
    "from openfermion.testing.testing_utils import random_quadratic_hamiltonian\n",
    "from openfermion.ops import QuadraticHamiltonian\n",
    "from openfermion.linalg.givens_rotations import givens_decomposition_square\n",
    "from openfermion.circuits.primitives import optimal_givens_decomposition\n",
    "\n",
    "from scipy.sparse import csc_matrix\n",
    "from scipy.linalg import expm\n",
    "from scipy.special import binom\n",
    "import scipy as sp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up the number of orbs and a random quadratic Hamiltonian\n",
    "norbs = 4\n",
    "ikappa = random_quadratic_hamiltonian(norbs, conserves_particle_number=True, real=True, expand_spin=False, seed=2)\n",
    "ikappa_mat = of.get_sparse_operator(of.get_fermion_operator(ikappa), n_qubits=norbs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up initial full Hilbert space wavefunction\n",
    "wfn = np.zeros((2**(norbs), 1), dtype=np.complex128)\n",
    "wfn[int(\"\".join([str(x) for x in [1] * (norbs//2) + [0] * (norbs//2)]), 2), 0] = 1  # alpha1-alpha2 Hartree-Fock state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# full space unitary\n",
    "time = np.random.random()\n",
    "hamiltonian_evolution = expm(-1j * ikappa_mat.todense() * time) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build a FQE operator corresponding to the Hamiltonian evolution\n",
    "fqe_ham = RestrictedHamiltonian((ikappa.n_body_tensors[1, 0],))\n",
    "assert fqe_ham.quadratic()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize the FQE wavefunction [n-elec, sz, norbs]\n",
    "fqe_wfn = fqe.Wavefunction([[norbs // 2, norbs//2, norbs]])\n",
    "hf_wf = np.zeros((int(binom(norbs, norbs // 2)), 1), dtype=np.complex128)\n",
    "hf_wf[0, 0] = 1  # right most bit is zero orbital.\n",
    "fqe_wfn.set_wfn(strategy='from_data',\n",
    "                raw_data={(norbs//2, norbs//2): hf_wf})\n",
    "fqe_wfn.print_wfn()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Aside\n",
    "\n",
    "The FQE and cirq/OpenFermion have reverse ordering of the occupation number vectors.  Give a bitstring representing occupations 01001 one commonly wants to determine the index of the orbitals used in the 2-orbital Slater determinant.  In OpenFermion the orbitals used are 1 and 4.  In FQE the orbitals used are 0 and 3.  The difference is that OpenFermion starts indexing from the left most bit whereas FQE starts indexing from the right most bit. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize wavefunction [n-elec, sz, norbs]\n",
    "fqe_wfn = fqe.Wavefunction([[norbs // 2, norbs//2, norbs]])\n",
    "hf_wf = np.zeros((int(binom(norbs, norbs // 2)), 1), dtype=np.complex128)\n",
    "hf_wf[0, 0] = 1  # right most bit is zero orbital. 2-alpha electron HF is 0011.  This is the first element of our particle conserved Hilbert space\n",
    "fqe_wfn.set_wfn(strategy='from_data',\n",
    "                raw_data={(norbs//2, norbs//2): hf_wf})\n",
    "fqe_wfn.print_wfn()\n",
    "\n",
    "# OpenFermion indexes from the left as zero.  The RHF state is 1100\n",
    "wfn[int(\"\".join([str(x) for x in [1] * (norbs//2) + [0] * (norbs//2)]), 2), 0] = 1\n",
    "\n",
    "print(\"\\nHF-Vector-OpenFermion \", bin(int(\"\".join([str(x) for x in [1] * (norbs//2) + [0] * (norbs//2)]), 2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fqe_wfn.print_wfn()  # before\n",
    "fqe_wfn = fqe_wfn.time_evolve(time, fqe_ham)\n",
    "fqe_wfn.print_wfn()  # after"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_wfn = hamiltonian_evolution @ wfn  # Full space time-evolution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare 1-RDMs\n",
    "\n",
    "It is easy to compare the 1-RDMs of the wavefunctions to confirm they are representing the same state.  We will compute the 1-RDM for the `wfn` and `fqe_wfn` and compare to the opdm computed from `fqe.to_cirq`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def get_opdm(wfn: Union[sp.sparse.coo_matrix, sp.sparse.csr_matrix,\n",
    "                        sp.sparse.csc_matrix], nso: int) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    A slow way to get 1-RDMs from density matrices or states\n",
    "    \"\"\"\n",
    "    if isinstance(wfn, (sp.sparse.coo_matrix, sp.sparse.csc_matrix)):\n",
    "        wfn = wfn.tocsr()\n",
    "        wfn = wfn.toarray()\n",
    "    opdm = np.zeros((nso, nso), dtype=wfn.dtype)\n",
    "    a_ops = [\n",
    "        of.get_sparse_operator(of.jordan_wigner(of.FermionOperator(((p, 0)))),\n",
    "                               n_qubits=nso)\n",
    "        for p in range(nso)\n",
    "    ]\n",
    "    for p, q in product(range(nso), repeat=2):\n",
    "        operator = a_ops[p].conj().T @ a_ops[q]\n",
    "        if wfn.shape[0] == wfn.shape[1]:  # we are working with a density matrix\n",
    "            val = np.trace(wfn @ operator)\n",
    "        else:\n",
    "            val = wfn.conj().transpose() @ operator @ wfn\n",
    "        # print((p, q, r, s), val.toarray()[0, 0])\n",
    "        if isinstance(val, (float, complex, int)):\n",
    "            opdm[p, q] = val\n",
    "        else:\n",
    "            opdm[p, q] = val[0, 0]\n",
    "\n",
    "    return opdm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fqe_opdm = fqe_wfn.rdm('i^ j') # get the FQE-opdm\n",
    "\n",
    "# get the exactly evolved opdm\n",
    "initial_opdm = np.diag([1] * (norbs//2) + [0] * (norbs//2))\n",
    "u = expm(1j * ikappa.n_body_tensors[1, 0] * time)\n",
    "final_opdm = u @ initial_opdm @ u.conj().T\n",
    "print(final_opdm)\n",
    "print()\n",
    "\n",
    "# contract the wavefunction opdm\n",
    "test_opdm = get_opdm(final_wfn, norbs)\n",
    "print(test_opdm)\n",
    "\n",
    "assert np.allclose(test_opdm.real, final_opdm.real)\n",
    "assert np.allclose(test_opdm.imag, final_opdm.imag)\n",
    "assert np.allclose(final_opdm, fqe_opdm)\n",
    "\n",
    "cirq_wfn = fqe.to_cirq(fqe_wfn).reshape((-1, 1))\n",
    "\n",
    "cfqe_opdm = get_opdm(cirq_wfn, 2 * norbs)\n",
    "print()\n",
    "print(cfqe_opdm[::2, ::2])\n",
    "assert np.allclose(final_opdm, cfqe_opdm[::2, ::2])\n",
    "\n",
    "# cirq evolution\n",
    "qubits = cirq.LineQubit.range(norbs)\n",
    "prep = cirq.Moment([cirq.X.on(qubits[0]), cirq.X.on(qubits[1])])\n",
    "rotation = cirq.Circuit(optimal_givens_decomposition(qubits, u.conj().copy()))\n",
    "circuit = prep + rotation\n",
    "\n",
    "final_state = circuit.final_state_vector().reshape((-1, 1))\n",
    "cirq_opdm = get_opdm(final_state, norbs)\n",
    "print()\n",
    "print(cirq_opdm)\n",
    "\n",
    "assert np.allclose(final_opdm, cirq_opdm)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Givens rotations with FQE\n",
    "\n",
    "We can use the same circuit generated by OpenFermion except apply the Givens rotations to the FQE wavefunction.\n",
    "Each Givens rotation gate sequence is mapped to \n",
    "$$\n",
    "e^{-i \\theta (a_{p}^{\\dagger}a_{q} - a_{q}^{\\dagger}a_{p})}e^{i \\phi a_{p}^{\\dagger}a_{p}}\n",
    "$$\n",
    "which can be applies sequentially to the state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the basis rotation circuit\n",
    "rotations, diagonal = givens_decomposition_square(u.conj().T)\n",
    "\n",
    "# Reinitialize the wavefunction to Hartree-Fock\n",
    "fqe_wfn = fqe.Wavefunction([[norbs // 2, norbs // 2, norbs]])\n",
    "hf_wf = np.zeros((int(binom(norbs, norbs // 2)), 1),\n",
    "                 dtype=np.complex128)\n",
    "hf_wf[0, 0] = 1  # right most bit is zero orbital.\n",
    "fqe_wfn.set_wfn(strategy='from_data',\n",
    "                raw_data={(norbs // 2, norbs // 2): hf_wf})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Iterate through each layer and time evolve by the appropriate fermion operators\n",
    "for layer in rotations:\n",
    "    for givens in layer:\n",
    "        i, j, theta, phi = givens\n",
    "        op = of.FermionOperator(((2 * j, 1), (2 * j, 0)), coefficient=-phi)\n",
    "        fqe_wfn = fqe_wfn.time_evolve(1.0, op)\n",
    "\n",
    "        op = of.FermionOperator(((2 * i, 1), (2 * j, 0)), coefficient=-1j * theta) + of.FermionOperator(((2 * j, 1), (2 * i, 0)), coefficient=1j * theta)\n",
    "        fqe_wfn = fqe_wfn.time_evolve(1.0, op)\n",
    "\n",
    "# evolve the last diagonal phases\n",
    "for idx, final_phase in enumerate(diagonal):\n",
    "    if not np.isclose(final_phase, 1.0):\n",
    "        op = of.FermionOperator(((2 * idx, 1), (2 * idx, 0)), -np.angle(final_phase))\n",
    "        fqe_wfn = fqe_wfn.time_evolve(1.0, op)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# map the wavefunction to a cirq wavefunction and get the opdm\n",
    "crq_wfn = fqe.to_cirq(fqe_wfn).reshape((-1, 1))\n",
    "crq_opdm = get_opdm(crq_wfn, 2 * norbs)\n",
    "print(crq_opdm[::2, ::2])\n",
    "\n",
    "# check if it is the same as our other 1-RDMs\n",
    "assert np.allclose(final_opdm, crq_opdm[::2, ::2])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
